#include <pcap.h>
#include <iostream>
#include <fstream>
#include <ctime>
#include <map>

#pragma comment(lib, "wpcap.lib")

// 定义统计结构
struct TrafficRecord {
    long long totalLen = 0;
};

std::ofstream logFile("network_log.csv");
std::map<std::string, TrafficRecord> srcStats;  // 源地址统计
std::map<std::string, TrafficRecord> dstStats;  // 目标地址统计
time_t lastStatTime = time(nullptr);            // 上次统计时间

// 生成时间戳
std::string getTimestamp() {
    time_t now = time(nullptr);
    char buf[20];
    strftime(buf, sizeof(buf), "%Y-%m-%d %H:%M:%S", localtime(&now));
    return std::string(buf);
}

// 处理数据包回调
void packetHandler(u_char* userData, const pcap_pkthdr* pkthdr, const u_char* packet) {
    // 解析以太网帧
    struct ether_header* ethHdr = (struct ether_header*)packet;
    if (ntohs(ethHdr->ether_type) != 0x0800) return; // 仅处理IPv4

    // 解析IP头部
    struct ip* ipHdr = (struct ip*)(packet + 14);
    char srcIP[INET_ADDRSTRLEN], dstIP[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(ipHdr->ip_src), srcIP, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &(ipHdr->ip_dst), dstIP, INET_ADDRSTRLEN);

    // 格式化MAC地址
    char srcMAC[18], dstMAC[18];
    sprintf(srcMAC, "%02X:%02X:%02X:%02X:%02X:%02X",
        ethHdr->ether_shost[0], ethHdr->ether_shost[1],
        ethHdr->ether_shost[2], ethHdr->ether_shost[3],
        ethHdr->ether_shost[4], ethHdr->ether_shost[5]);
    sprintf(dstMAC, "%02X:%02X:%02X:%02X:%02X:%02X",
        ethHdr->ether_dhost[0], ethHdr->ether_dhost[1],
        ethHdr->ether_dhost[2], ethHdr->ether_dhost[3],
        ethHdr->ether_dhost[4], ethHdr->ether_dhost[5]);

    // 写入日志
    logFile << getTimestamp() << "," 
            << srcMAC << "," << srcIP << ","
            << dstMAC << "," << dstIP << ","
            << pkthdr->len << "\n";

    // 更新统计
    std::string srcKey = std::string(srcMAC) + "|" + srcIP;
    std::string dstKey = std::string(dstMAC) + "|" + dstIP;
    srcStats[srcKey].totalLen += pkthdr->len;
    dstStats[dstKey].totalLen += pkthdr->len;

    // 每分钟输出统计
    time_t currentTime = time(nullptr);
    if (difftime(currentTime, lastStatTime) >= 60) {
        std::ofstream statFile("network_stats.csv", std::ios::app);
        statFile << "统计时间,统计类型,地址,总数据长度\n";
        for (auto& entry : srcStats) {
            statFile << getTimestamp() << ",源MAC," << entry.first << "," << entry.second.totalLen << "\n";
        }
        for (auto& entry : dstStats) {
            statFile << getTimestamp() << ",目标MAC," << entry.first << "," << entry.second.totalLen << "\n";
        }
        srcStats.clear();
        dstStats.clear();
        lastStatTime = currentTime;
    }
}

int main() {
    pcap_if_t* allDevs;
    char errBuf[PCAP_ERRBUF_SIZE];

    // 获取网卡列表
    if (pcap_findalldevs(&allDevs, errBuf) == -1) {
        std::cerr << "Error: " << errBuf << std::endl;
        return 1;
    }

    // 选择第一个网卡
    pcap_t* handle = pcap_open_live(allDevs->name, 65536, 1, 1000, errBuf);
    if (!handle) {
        std::cerr << "无法打开网卡: " << errBuf << std::endl;
        return 1;
    }

    // 设置过滤规则（仅IPv4）
    struct bpf_program filter;
    pcap_compile(handle, &filter, "ip", 1, 0);
    pcap_setfilter(handle, &filter);

    // 写入CSV表头
    logFile << "时间,源MAC,源IP,目标MAC,目标IP,帧长度\n";

    // 开始抓包
    pcap_loop(handle, 0, packetHandler, nullptr);

    pcap_close(handle);
    return 0;
}